// Copyright 2022 Chen Jun

#include "hnurm_armor/extended_kalman_filter.hpp"

#include <utility>

namespace hnurm
{
ExtendedKalmanFilter::ExtendedKalmanFilter(
    VecVecFunc             f,
    VecVecFunc             h,
    VecMatFunc             j_f,
    VecMatFunc             j_h,
    VoidMatFunc            u_q,
    VecMatFunc             u_r,
    const Eigen::MatrixXd &P0
)
    : f(std::move(f)),
      h(std::move(h)),
      jacobian_f(std::move(j_f)),
      jacobian_h(std::move(j_h)),
      update_Q(std::move(u_q)),
      update_R(std::move(u_r)),
      P_post(P0),
      n((int)P0.rows()),
      I(Eigen::MatrixXd::Identity(n, n)),
      x_pri(n),
      x_post(n)
{
}

void ExtendedKalmanFilter::setState(const Eigen::VectorXd &x0)
{
    x_post = x0;
}

Eigen::MatrixXd ExtendedKalmanFilter::predict()
{
    F = jacobian_f(x_post), Q = update_Q();

    x_pri = f(x_post);
    P_pri = F * P_post * F.transpose() + Q;

    // handle the case when there will be no measurement before the next predict
    x_post = x_pri;
    P_post = P_pri;

    return x_pri;
}

Eigen::MatrixXd ExtendedKalmanFilter::update(const Eigen::VectorXd &z)
{
    H = jacobian_h(x_pri), R = update_R(z);

    K      = P_pri * H.transpose() * (H * P_pri * H.transpose() + R).inverse();
    x_post = x_pri + K * (z - h(x_pri));
    P_post = (I - K * H) * P_pri;

    return x_post;
}
}  // namespace hnurm
